SoftmaxCrossEntropyWithLogits

计算 Softmax 交叉熵损失及梯度。

该算子首先对输入 logits 进行 Softmax 归一化得到概率，然后计算其与 labels 的交叉熵损失。如果启用了 need_grads，则会计算损失相对于 logits 的梯度。

算法逻辑如下：

  1. Softmax:

\[p_{i,j} = \frac{e^{x_{i,j}}}{\sum_{k} e^{x_{i,k}}}\]

  1. Cross Entropy Loss:

\[loss_i = - \sum_{j} y_{i,j} \log(p_{i,j})\]

  1. Gradients (当 need_grads=1 时):

\[\frac{\partial loss}{\partial x_{i,j}} = p_{i,j} - y_{i,j}\]
输入：

  • logits - 输入数据地址（未归一化的对数概率）。形状为 \([batch\_size, num\_of\_classes]\)。

  • labels - 标签数据地址。形状为 \([batch\_size, num\_of\_classes]\)。

  • probs - 输出概率地址（Softmax 结果）。

  • grads - 输出梯度地址。如果 need_grads 为 0，可忽略。

  • output - 输出损失值地址（通常为标量或每个样本的损失）。

  • sum_data - 中间计算缓冲区（Workspace），用于存储行和等临时数据。

  • batch_size - 批大小。

  • num_of_classes - 类别数量。

  • need_grads - 是否需要计算梯度 (0: 不计算, 1: 计算)。

  • core_mask - 核掩码（仅适用于共享存储版本）。

输出：

  • probs - 更新后的概率分布。

  • grads - 计算得到的梯度（如启用）。

  • output - 计算得到的交叉熵损失。

支持平台：

FT78NE MT7004

备注

  • FT78NE 支持 Fp32 和 Int8 数据类型。

  • MT7004 支持 Fp32 和 FP16 数据类型。

  • 输入数据必须是二维矩阵，行优先存储。

  • Int8 版本通常采用混合精度计算：输入/标签为 int8_t，但输出（概率、梯度、损失）为 float 以保证精度。

共享存储版本:

void fp_softmax_cross_entropy_with_logits_s(float *logits, float *labels, float *probs, float *grads, float *output, float *sum_data, int batch_size, int num_of_classes, int need_grads, int core_mask)
void hp_softmax_cross_entropy_with_logits_s(float16 *logits, float16 *labels, float16 *probs, float16 *grads, float16 *output, float16 *sum_data, int batch_size, int num_of_classes, int need_grads, int core_mask)
void i8_softmax_cross_entropy_with_logits_s(int8_t *logits, int8_t *labels, float *probs, float *grads, float *output, float *sum_data, int batch_size, int num_of_classes, int need_grads, int core_mask)

C调用示例（FT78NE - Int8）：

 1#include <stdio.h>
 2#include <stdint.h>
 3
 4int main(int argc, char* argv[]) {
 5    // 假设所有数据均位于DDR空间
 6    int8_t* logits = (int8_t*)0xC0000000;
 7    int8_t* labels = (int8_t*)0xC1000000;
 8    float* probs   = (float*)0xC2000000;
 9    float* grads   = (float*)0xC3000000;
10    float* output  = (float*)0xC4000000;
11    float* sum_data= (float*)0xC5000000; // Workspace
12
13    int batch_size = 128;
14    int num_of_classes = 1000;
15    int need_grads = 1;
16    int core_mask = 0xff; // 使用所有核心
17
18    // i8版本：输入为int8，输出及中间计算为float
19    i8_softmax_cross_entropy_with_logits_s(logits, labels, probs, grads, output,
20                                           sum_data, batch_size, num_of_classes,
21                                           need_grads, core_mask);
22
23    return 0;
24}

私有存储版本:

void fp_softmax_cross_entropy_with_logits_p(float *logits, float *labels, float *probs, float *grads, float *output, float *sum_data, int batch_size, int num_of_classes, int need_grads)
void hp_softmax_cross_entropy_with_logits_p(float16 *logits, float16 *labels, float16 *probs, float16 *grads, float16 *output, float16 *sum_data, int batch_size, int num_of_classes, int need_grads)
void i8_softmax_cross_entropy_with_logits_p(int8_t *logits, int8_t *labels, float *probs, float *grads, float *output, float *sum_data, int batch_size, int num_of_classes, int need_grads)

C调用示例（MT7004 - FP16）：

 1#include <stdio.h>
 2
 3int main(int argc, char* argv[]) {
 4    // 假设所有数据均位于AM空间
 5    float16* logits = (float16*)0x10010000;
 6    float16* labels = (float16*)0x10020000;
 7    float16* probs  = (float16*)0x10030000;
 8    float16* grads  = (float16*)0x10040000;
 9    float16* output = (float16*)0x10050000;
10    float16* sum_data=(float16*)0x10060000;
11
12    int batch_size = 32;
13    int num_of_classes = 512;
14    int need_grads = 1;
15
16    hp_softmax_cross_entropy_with_logits_p(logits, labels, probs, grads, output,
17                                            sum_data, batch_size, num_of_classes,
18                                            need_grads);
19
20    return 0;
21}